// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/exec/assign_uniqueId.h>

#include <algorithm>
#include <utility>

namespace kumo::pollux::exec {
    AssignUniqueId::AssignUniqueId(
        int32_t operatorId,
        DriverCtx *driverCtx,
        const std::shared_ptr<const core::AssignUniqueIdNode> &planNode,
        int32_t uniqueTaskId,
        std::shared_ptr<std::atomic_int64_t> rowIdPool)
        : Operator(
              driverCtx,
              planNode->outputType(),
              operatorId,
              planNode->id(),
              "AssignUniqueId"),
          rowIdPool_(std::move(rowIdPool)) {
        POLLUX_USER_CHECK_LT(
            uniqueTaskId,
            kTaskUniqueIdLimit,
            "Unique 24-bit ID specified for AssignUniqueId exceeds the limit");
        uniqueValueMask_ = static_cast<int64_t>(uniqueTaskId) << 40;

        const auto numColumns = planNode->outputType()->size();
        identityProjections_.reserve(numColumns - 1);
        for (column_index_t i = 0; i < numColumns - 1; ++i) {
            identityProjections_.emplace_back(i, i);
        }

        resultProjections_.emplace_back(0, numColumns - 1);
        results_.resize(1);

        rowIdCounter_ = 0;
        maxRowIdCounterValue_ = 0;
    }

    void AssignUniqueId::addInput(RowVectorPtr input) {
        auto numInput = input->size();
        POLLUX_CHECK_NE(
            numInput, 0, "AssignUniqueId::addInput received empty set of rows");
        input_ = std::move(input);
    }

    RowVectorPtr AssignUniqueId::getOutput() {
        if (input_ == nullptr) {
            return nullptr;
        }
        generateIdColumn(input_->size());
        auto output = fillOutput(input_->size(), nullptr);
        input_ = nullptr;
        return output;
    }

    bool AssignUniqueId::isFinished() {
        return noMoreInput_ && input_ == nullptr;
    }

    void AssignUniqueId::generateIdColumn(vector_size_t size) {
        // Re-use memory for the ID vector if possible.
        VectorPtr &result = results_[0];
        if (result && result.use_count() == 1) {
            BaseVector::prepare_for_reuse(result, size);
        } else {
            result = BaseVector::create(BIGINT(), size, pool());
        }

        auto rawResults =
                result->as_unchecked<FlatVector<int64_t> >()->mutableRawValues();

        vector_size_t start = 0;
        while (start < size) {
            if (rowIdCounter_ >= maxRowIdCounterValue_) {
                requestRowIds();
            }

            const auto numAvailableIds =
                    std::min(maxRowIdCounterValue_ - rowIdCounter_, kRowIdsPerRequest);
            const vector_size_t end =
                    std::min(static_cast<int64_t>(size), start + numAvailableIds);
            POLLUX_CHECK_EQ(
                (rowIdCounter_ + (end - start)) & uniqueValueMask_,
                0,
                "Ran out of unique IDs at {}. Need {} more.",
                rowIdCounter_,
                (end - start));
            std::iota(
                rawResults + start, rawResults + end, uniqueValueMask_ | rowIdCounter_);

            rowIdCounter_ += (end - start);
            start = end;
        }
    }

    void AssignUniqueId::requestRowIds() {
        rowIdCounter_ = rowIdPool_->fetch_add(kRowIdsPerRequest);
        maxRowIdCounterValue_ =
                std::min(rowIdCounter_ + kRowIdsPerRequest, kMaxRowId);
    }
} // namespace kumo::pollux::exec
